package alloy.transform.atomize;

import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

import alloy.ast.ASTDepthFirstVisitor;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.Decls;
import alloy.ast.EmptyFormula;
import alloy.ast.EmptySetExpr;
import alloy.ast.FindCommand;
import alloy.ast.IdentityExpr;
import alloy.ast.Leaf;
import alloy.ast.LeafId;
import alloy.ast.LiteralIntExpr;
import alloy.ast.MultiplicityExpr;
import alloy.ast.Op;
import alloy.ast.QuantifiedFormula;
import alloy.ast.SigExpr;
import alloy.ast.Signature;
import alloy.ast.Signatures;
import alloy.ast.SumIntExpr;
import alloy.ast.UniversalExpr;
import alloy.ast.VarCreator;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.semantic.LocalScope;
import alloy.transform.LeafIdMap;
import alloy.type.SigType;
import alloy.util.Dbg;

/**
 * Stores state information of an AtomizationVisitor so that we can always set
 * up an instance of the visitor to resume atomization of a command. This
 * functionality is needed by the evaluator.
 * 
 * also stores the mapping required for a reverse name mapping to performed
 * after atomization.
 * 
 * This was previously known as AtomizationInfo
 */
class CommandState {

	/** set of SigTypes that are selected as atomic */
	final Set _atomicTypes = new LinkedHashSet();

	/** mapping from supertype to set of atomic types that partition it */
	final Map _superToAtomic = new LinkedHashMap();

	/** mapping from subtype its unique atomic supertype */
	final Map _subToAtomic = new LinkedHashMap();

	/**
	 * mapping of LeafId's for SigExpr's to their SigmaExpr replacements (in
	 * place of SigExpr's for these SigTypes), if applicable
	 */
	final Map _sigIdToSigmaExpr = new LinkedHashMap();

	/**
	 * mapping of LeafId's for fields to their SigmaExpr replacements, if
	 * applicable
	 */
	final Map _fieldIdToSigmaExpr = new LinkedHashMap();

	/**
	 * mapping of LeafId's for VariableExpr's to their SigmaExpr replacements
	 * currently used exclusively for skolem constants.
	 */
	final Map _varIdToSigmaExpr = new LinkedHashMap();

	/**
	 * Maps LeafId's of expressions (both in the atomized and original asts)
	 * back to their original forms. Expressions that got broken down in
	 * atomization will map to themselves.
	 * 
	 * NOTE: there is a slight problem with using a one-to-one mapping here.
	 * This is when signatures get partitioned (as in the case for an abstract
	 * type X that is extended by A, which was scoped to be atomic). In this
	 * case, you want the atoms of A to appear both in A and X, but a one-to-one
	 * mapping will only add the atoms to the SigExpr for A, leaving X empty.
	 * This is the direct cause for the empty univ bug.
	 * 
	 * The solution for this problem is a small implementation detail in
	 * alloy.api.Instance that'll automatically fill the supertypes. This may
	 * not sound very general, but indeed, it is clean, and the generality is
	 * uncalled for because for all other expressions that get
	 * atomized--skolems, quantified variables, and field, at most one level of
	 * partitioning occurs, and the mapping is correctly made to the original
	 * expressions in these cases (because there is no need to also map to the
	 * new expression, e.g. foo$A$B, as there is a need to map to the subtype
	 * expression, e.g. A, for signatures).
	 */
	public final Map _atomizedIdToOldLeafExpr = new LinkedHashMap();

	/**
	 * a mapping from SigType's to their TypeScope's (if any) used to determine
	 * if it was exact
	 */
	final Map _sigTypeToTypeScope = new LinkedHashMap();

	/**
	 * visitor that can be run to ensure that all LeafId's are set and
	 * consistent
	 */
	public final FixLeafIdsVisitor _fliv;

	final FindCommand _command;

	// private AtomizationInfo(final AtomizationInfo other) {
	// _allSigs = other._allSigs;
	// _atomizedIdToOldLeafExpr = other._atomizedIdToOldLeafExpr;
	// _atomicTypes = other._atomicTypes;
	// _superToAtomic = other._superToAtomic;
	// _subToAtomic = other._subToAtomic;
	// _sigTypeToTypeScope = other._sigTypeToTypeScope;
	// _sigIdToSigmaExpr = other._sigIdToSigmaExpr;
	// _fieldIdToSigmaExpr = other._fieldIdToSigmaExpr;
	// _varIdToSigmaExpr = other._varIdToSigmaExpr;
	// _curModuleInfo = other._curModuleInfo;
	// _fliv = other._fliv;
	// _createdStringToType = other._createdStringToType;
	// }

	CommandState(final FindCommand command_) {
		_command = command_;

		// final FixLeafIdsVisitor
		_fliv = new FixLeafIdsVisitor(_command.leafIdToMult);
		_fliv.setNonExprLeafIdMap(_command.sliv.getNonExprLeafIdMap());
		_fliv.setEmptyFormulaLeafId(_command.sliv.getEmptyFormulaLeafId());

		_command.stringToType = new LinkedHashMap(SigType.getStringToType());

	}

	/**
	 * this is a modified version of alloy.transform.SetLeafIdsVisitor. In
	 * particular, it handles skolem constants differently (doesn't do anything
	 * to them), and there're minor differences in the way VariableExpr's are
	 * handled in general.
	 */
	class FixLeafIdsVisitor extends ASTDepthFirstVisitor {

		/** current leaf */
		private LeafIdMap _curLeafIdMap;

		/** The id for an {@link EmptyFormula}, if we have seen one. */
		private LeafId _emptyFormulaId;

		/** current map from leaf ids to mult exprs */
		private Map _leafIdToMult;

		/**
		 * a mapping from the String representations of non-expression leaves to
		 * their leaf ids, for the purpose of reusing the ids. {@link Map} from
		 * {@link String} to {@link LeafId}.
		 */
		private Map _nonExprLeafIdMap = new LinkedHashMap();

		/**
		 * used to set the nonExprLeafIdMap. This is available so that we can
		 * use the original values before atomization.
		 * 
		 * @param map_
		 */
		void setNonExprLeafIdMap(Map map_) {
			_nonExprLeafIdMap = map_;
		}

		void setEmptyFormulaLeafId(LeafId id_) {
			_emptyFormulaId = id_;
		}

		public FixLeafIdsVisitor(Map leafIdToMult_) {
			_leafIdToMult = leafIdToMult_;
		}

		private LeafIdMap _genLeafIdMapping(Decls decls_, LeafIdMap parent_) {
			LeafIdMap map = new LeafIdMap(parent_);
			Iterator declIter = decls_.getDeclIter();
			while (declIter.hasNext()) {
				Decl decl = (Decl) declIter.next();
				Iterator vars = decl.getVariables().getVariableIter();
				while (vars.hasNext()) {
					Variable var = (Variable) vars.next();
					// System.out.println("generating id for " +
					// var.orignodeString());
					map.put(var.nodeString(), new LeafId());
				}
			}
			return map;
		}

		private void _handleNonExprLeaf(Leaf leaf_) {
			if (leaf_.getLeafId() == null) {
				LeafId id = null;
				// see if id has already been created
				String leafKey = leaf_.getClass().toString() + "##"
						+ leaf_.nodeString();
				if ((id = (LeafId) _nonExprLeafIdMap.get(leafKey)) == null) {
					// create new id if necessary
					_nonExprLeafIdMap.put(leafKey, id = new LeafId());
				}
				// set leaf id
				leaf_.setLeafId(id);
			}
		}

		private void _updateLeafIdToMult(Decls decls_) {
			if (_leafIdToMult != null) {
				Iterator declIter = decls_.getDeclIter();
				while (declIter.hasNext()) {
					Decl decl = (Decl) declIter.next();
					MultiplicityExpr multExpr = decl.getExpr();
					Iterator varIter = decl.getVariables().getVariableIter();
					while (varIter.hasNext()) {
						LeafId id = _curLeafIdMap
								.lookupLeafId(((Variable) varIter.next())
										.nodeString());
						if (!_leafIdToMult.containsKey(id)) {
							_leafIdToMult.put(id, multExpr.copy());
						}
					}
				}
			}
		}

		public void visit(ComprehensionExpr compExpr_) {
			visit((VarCreator) compExpr_);
		}

		public void visit(EmptyFormula emptyFormula_) {
			if (_emptyFormulaId == null)
				_emptyFormulaId = new LeafId();
			emptyFormula_.setLeafId(_emptyFormulaId);
		}

		public void visit(EmptySetExpr emptySetExpr_) {
			_handleNonExprLeaf(emptySetExpr_);
		}

		public void visit(IdentityExpr idenExpr_) {
			_handleNonExprLeaf(idenExpr_);
		}

		public void visit(LiteralIntExpr litIntExpr_) {
			_handleNonExprLeaf(litIntExpr_);
		}

		public void visit(Op op_) {
			_handleNonExprLeaf(op_);
		}

		public void visit(QuantifiedFormula quantForm_) {
			visit((VarCreator) quantForm_);
		}

		public void visit(SigExpr sigExpr_) {
			// Dbg.info(sigExpr_.replacednodeString());
			if (sigExpr_.getLeafId() == null) {
				// get signature object
				Signature sig = _command.moduleInfo.getSignature(sigExpr_
						.getSig(), null);

				// if the signature is not in this module, try to find it in
				// other modules
				// this usually doesn't happen. but a bug report has triggered
				// this case.
				// normally, curModuleInfo will include all the signatures (even
				// those
				// in imported modules). need to find out the exact rules.
				if (sig == null) {
					for (Iterator signaturesIter = AtomizationState.state._allSigs
							.iterator(); signaturesIter.hasNext();) {
						Signatures sigs = (Signatures) signaturesIter.next();
						for (Iterator sigIter = sigs.getSignatureIter(); sigIter
								.hasNext();) {
							Signature curSig = (Signature) sigIter.next();
							if (curSig.getName().equals(sigExpr_.getSig())) {
								sig = curSig;
							}
						}
					}

				}

				Dbg.chk(sig != null, "null signature for "
						+ sigExpr_.nodeString());
				// get id
				LeafId id = sig.getLeafId();
				Dbg.chk(id != null, "null signature id for "
						+ sigExpr_.nodeString());
				sigExpr_.setLeafId(id);

			}
		}

		public void visit(SumIntExpr sumIntExpr_) {
			visit((VarCreator) sumIntExpr_);
		}

		public void visit(UniversalExpr univExpr_) {
			_handleNonExprLeaf(univExpr_);
		}

		// nodes which can create new names

		public void visit(VarCreator creator_) {
			// save old leaf id map
			LeafIdMap oldLeafIdMap = _curLeafIdMap;
			if ((_curLeafIdMap = creator_.getLeafIdMap()) == null) {
				// generate leaf ids for newly bound names
				LocalScope localScope = creator_.getLocalScope();
				if (localScope != null) {
					_curLeafIdMap = localScope
							.generateLeafIdMapping(oldLeafIdMap);
				} else {
					// generate id mapping from decls
					_curLeafIdMap = _genLeafIdMapping(creator_.getDecls(),
							oldLeafIdMap);
				}
				creator_.setLeafIdMap(_curLeafIdMap);
			}

			_updateLeafIdToMult(creator_.getDecls());

			// visit all children
			visit((alloy.ast.Node) creator_);
			// restore leaf id map
			_curLeafIdMap = oldLeafIdMap;
		}

		public void visit(VariableExpr variableExpr_) {

			// see if variable is bound in current mapping
			String varString = variableExpr_.getVar().getId().nodeString();
			LeafId id = (_curLeafIdMap == null) ? null : _curLeafIdMap
					.lookupLeafId(varString);
			if (id == null) { // not a quantified variable
				// variable is either a skolemized variable or a field
				// both of these should've been taken off as we created them

				Dbg.chk(variableExpr_.getLeafId() != null,
						"problem with FixLeafIdsVisitor.  Non-quantified variable has null id "
								+ variableExpr_);

				// get signature object
				// Signature sig = _getSignatureObject(varString,
				// variableExpr_.getType());
				// Dbg.info(sig.getTypedName());
				// id = sig.getFieldLeafId(varString);
			} else {
				// quantified variable, so set flag
				variableExpr_.isQuantified = true;
				// set the quantified variable's LeafId that we found in the map
				// don't try to put in "optimizations" to see whether or not the
				// id's already
				// been set. a crucial purpose of this visitor was to correct
				// LeafId
				// inconsistencies caused by copying from a single original
				// VarCreator.
				variableExpr_.setLeafId(id);
				Dbg.chk(id != null, "no id for " + varString + " at "
						+ variableExpr_.getLocation().toString());
				variableExpr_.setLeafId(id);
			}

		}

	}
}
